# coding=utf-8

"""
    Created by Kalyter on 2017-11-22.
"""
import matplotlib.pyplot as plt
import matplotlib.cm as cm
from matplotlib.colors import LogNorm
import numpy as np


def complex():
    x, y = np.random.rand(10), np.random.rand(10)
    z = (np.random.rand(90000) + np.linspace(0, 1, 90000)).reshape(30, 30)
    plt.imshow(z+10, extent=(np.amin(x), np.amax(x), np.amin(y), np.amax(y)),
            cmap=cm.hot, norm=LogNorm())
    plt.colorbar()


def simple():
    x = [[1, 2], [3, 4], [5, 6]]
    plt.imshow(x)
    plt.show()


def color_print():
    X = [[0, 0.25], [0.5, 0.75]]

    fig = plt.figure()
    ax = fig.add_subplot(121)
    im = ax.imshow(X, cmap=plt.get_cmap('hot'))
    plt.colorbar(im, shrink=0.5)

    ax = fig.add_subplot(122)
    im = ax.imshow(X, cmap=plt.get_cmap('hot'), interpolation='nearest',
                   vmin=0, vmax=1)
    plt.colorbar(im, shrink=0.2)
    plt.show()


def draw_heatmap(data, xlabels, ylabels):
    # cmap=cm.Blues
    cmap = cm.get_cmap('rainbow', 1000)
    figure = plt.figure(facecolor='w')
    ax = figure.add_subplot(1, 1, 1, position=[0.1, 0.15, 0.8, 0.8])
    ax.set_yticks(range(len(ylabels)))
    ax.set_yticklabels(ylabels)
    ax.set_xticks(range(len(xlabels)))
    ax.set_xticklabels(xlabels)
    vmax = data[0][0]
    vmin = data[0][0]
    for i in data:
        for j in i:
            if j > vmax:
                vmax = j
            if j < vmin:
                vmin = j
    map = ax.imshow(data, interpolation='nearest', cmap=cmap, aspect='auto',
                    vmin=vmin, vmax=vmax)
    cb = plt.colorbar(mappable=map, cax=None, ax=None, shrink=0.5)
    plt.show()


if __name__ == '__main__':
    a = np.random.rand(10, 10)
    print(a)
    xlabels = ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J']
    ylabels = ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j']
    draw_heatmap(a, xlabels, ylabels)
